#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : RAG-demo
# @File    : app.py
# @IDE     : PyCharm
# @Author  :ZH
# @Time    : 2025/6/27 16:22
import os
from io import BytesIO

from fastapi import FastAPI, Body, UploadFile, File
from util import SQLiteDB, calculate_md5, ZhipuFileUtil, easy_split_content, bce_embedding_server, KnowledgeMilvus, VLLMFileUtil
from util.elasticsearch_util import ElasticsearchBM25Client
from util.logger_setup_util import LoggerSetup

logger = LoggerSetup(__name__).get_logger()

# 获取环境变量
ZHIPU_API_KEY = os.getenv("ZHIPU_API_KEY", "06e1eae173bf9fe678ca21dd32cf0fa3.oJn4E7me2OOB084p")
MILVUS_URI = os.getenv("MILVUS_URI", "http://123.57.164.1:19530/")
KNOWLEDGE_COLLECTION_NAME = os.getenv("KNOWLEDGE_COLLECTION_NAME", "knowledge_collection")
ELASTICSEARCH_HOST = os.getenv("ELASTICSEARCH_HOST", "elasticsearch.vip.cpolar.cn")
ELASTICSEARCH_PORT = int(os.getenv("ELASTICSEARCH_PORT", 80))
ELASTICSEARCH_INDEX_NAME = os.getenv("ELASTICSEARCH_INDEX_NAME", "bm25-index")

# 初始化一些对象
elasticsearch_object = ElasticsearchBM25Client(host=ELASTICSEARCH_HOST, port=ELASTICSEARCH_PORT,
                                               index_name="bm25-index")
knowledge_milvus = KnowledgeMilvus(uri=MILVUS_URI)

# 初始化SQLite数据库
db_path = "db/rag_demo.db"
sqlite_db = SQLiteDB(db_path=db_path)
app = FastAPI(title="RAG Demo API")

# 临时使用本地数据库SQLite
# 检测数据库是否存在
table_name = {"knowledge", "knowledge_file", "dialogue", "file_summary_page", "file_summary"}
create_table_sql = {
    "knowledge": """
        CREATE TABLE IF NOT EXISTS knowledge (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT NOT NULL,
            description TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """,
    "knowledge_file": """
        CREATE TABLE IF NOT EXISTS knowledge_file (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            knowledge_id INTEGER NOT NULL,
            file_path TEXT NOT NULL,
            file_key TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (knowledge_id) REFERENCES knowledge(id)
        )
    """,
    "dialogue": """
        CREATE TABLE IF NOT EXISTS dialogue (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            human_message TEXT NOT NULL,
            ai_message TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """,
    "file_summary_page": """
    CREATE TABLE IF NOT EXISTS file_summary_page (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    knowledge_id INTEGER NOT NULL,
    file_key TEXT NOT NULL,
    page_number INTEGER NOT NULL,
    summary TEXT NOT NULL,
    )
    """,
    "file_summary": """
    CREATE TABLE IF NOT EXISTS file_summary (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    knowledge_id INTEGER NOT NULL,
    file_key TEXT NOT NULL,
    summary TEXT NOT NULL,
    )
    """
}

table_title = {
    "knowledge": ["id", "name", "description", "created_at"],
    "knowledge_file": ["id", "knowledge_id", "file_path", "file_key", "created_at"],
    "dialogue": ["id", "human_message", "ai_message", "created_at"],
    "file_summary_page": ["id", "knowledge_id", "file_key", "page_number", "summary"],
    "file_summary": ["id", "knowledge_id", "file_key", "summary"]
}

with sqlite_db as db_client:
    for table in table_name:
        if not db_client.table_exists(table):
            db_client.create_table(create_table_sql[table])
            logger.info(f"Table '{table}' created successfully.")
        else:
            logger.info(f"Table '{table}' already exists.")

# 初始化ES
elasticsearch_object.create_index()


@app.get("/knowledge/list", summary="获取知识库列表")
async def get_knowledge_list():
    """
    获取知识库列表
    """
    sql = "SELECT id, name, description, created_at FROM knowledge"
    try:
        with sqlite_db as db_client:
            result = db_client.execute(sql)
        return {
            "status": "success",
            "data": [dict(zip(table_title.get('knowledge'), row)) for row in result] if table_title.get(
                'knowledge') and result else []
        }
    except Exception as e:
        logger.error(f"Error fetching knowledge list: {e}")
        return {"status": "error", "message": str(e)}


@app.post("/knowledge/add", summary="添加知识库")
async def add_knowledge(
        name: str = Body(..., description="知识库名称"),
        description: str = Body(None, description="知识库描述")
):
    """
    添加知识库
    """
    sql = "INSERT INTO knowledge (name, description) VALUES (?, ?)"
    try:
        with sqlite_db as db_client:
            db_client.execute(sql, (name, description))
        return {"status": "success", "message": "Knowledge added successfully"}
    except Exception as e:
        logger.error(f"Error adding knowledge: {e}")
        return {"status": "error", "message": str(e)}


@app.post('/knowledge/file/add', summary="添加知识库文件")
async def add_knowledge_file(
        knowledge_id: int,
        file: UploadFile = File(...)
):
    """
    添加知识库文件,Demo阶段就使用同步方案来处理文件
    """
    # 获取文件的md5
    file_key = await calculate_md5(file=file)
    # 查询当前知识库中是否存在相同的文件
    sql = "SELECT id FROM knowledge_file WHERE knowledge_id = ? AND file_key = ?"
    try:
        with sqlite_db as db_client:
            existing_file = db_client.execute(sql, (knowledge_id, file_key))
            if existing_file:
                return {"status": "error", "message": "File already exists in the knowledge base"}
    except Exception as e:
        logger.error(f"Error checking existing file: {e}")
        return {"status": "error", "message": str(e)}
    # 文件留存到本地
    file_path = f"knowledge_files/{knowledge_id}/{file.filename}"
    os.makedirs(os.path.dirname(file_path), exist_ok=True)
    with open(file_path, "wb") as f:
        content = await file.read()
        f.write(content)
    sql = "INSERT INTO knowledge_file (knowledge_id, file_path, file_key) VALUES (?, ?, ?)"
    try:
        with sqlite_db as db_client:
            db_client.execute(sql, (knowledge_id, file_path, file_key))
    except Exception as e:
        logger.error(f"Error adding knowledge file: {e}")
        return {"status": "error", "message": str(e)}
    # 提取文本内容
    zhipu_file_util = ZhipuFileUtil(api_key=ZHIPU_API_KEY)
    file_content = await zhipu_file_util.extract_file_content(file_object=(file.filename, BytesIO(content)))
    # 文本数据切分
    file_split = easy_split_content(file_content, chunk_size=512)
    # 保存切分后的文本转成向量 [{"content": "文本内容", "embedding": [向量数据]}]
    content_embedding = await bce_embedding_server(content=file_split)
    knowledge_milvus.create_collection(collection_name=KNOWLEDGE_COLLECTION_NAME, dim=768, description="知识库向量存储")
    knowledge_milvus.load_collection(collection_name=KNOWLEDGE_COLLECTION_NAME)
    save_list = [
        {
            "knowledge_id": knowledge_id,
            "file_name": file.filename,
            "file_hase": file_key,
            "document": item['content'],
            "vector": item['embedding']
        } for item in content_embedding
    ]
    knowledge_milvus.insert(
        collection_name=KNOWLEDGE_COLLECTION_NAME,
        data=save_list
    )
    # 存储到ES 分批插入
    for index in range(0, len(save_list), 100):
        batch = save_list[index:index + 100]
        try:
            elasticsearch_object.add_documents(batch)
            logger.info(f"Added {len(batch)} documents to Elasticsearch index '{ELASTICSEARCH_INDEX_NAME}'.")
        except Exception as e:
            logger.error(f"Error adding documents to Elasticsearch: {e}")
            return {"status": "error", "message": str(e)}

    return {
        "status": "success",
        "message": "File added to knowledge base successfully",
        "file_content": file_content
    }


@app.delete('/knowledge/file/delete', summary="删除知识库文件")
async def delete_knowledge_file(
        knowledge_id: int,
        file_key: str
):
    """
    删除知识库文件
    """
    sql = "DELETE FROM knowledge_file WHERE knowledge_id = ? AND file_key = ?"
    try:
        with sqlite_db as db_client:
            db_client.execute(sql, (knowledge_id, file_key))
        # 删除文件
        file_path = f"knowledge_files/{knowledge_id}/{file_key}"
        if os.path.exists(file_path):
            os.remove(file_path)
        # 删除Milvus中的向量数据
        knowledge_milvus.delete(
            collection_name=KNOWLEDGE_COLLECTION_NAME,
            filter=f"knowledge_id = {knowledge_id} AND file_hase = '{file_key}'"
        )
        # 删除ES中的文档
        elasticsearch_object.delete_document(
            knowledge_id=knowledge_id,
            file_hase=file_key
        )
        return {"status": "success", "message": "File deleted successfully"}
    except Exception as e:
        logger.error(f"Error deleting knowledge file: {e}")
        return {"status": "error", "message": str(e)}


@app.delete('/knowledge/delete', summary="删除知识库")
async def delete_knowledge(
        knowledge_id: int
):
    """
    删除知识库
    """
    sql = "DELETE FROM knowledge WHERE id = ?"
    try:
        with sqlite_db as db_client:
            db_client.execute(sql, (knowledge_id,))
        # 删除相关的文件记录
        sql_files = "DELETE FROM knowledge_file WHERE knowledge_id = ?"
        with sqlite_db as db_client:
            db_client.execute(sql_files, (knowledge_id,))
        # 删除Milvus中的向量数据
        knowledge_milvus.delete(
            collection_name=KNOWLEDGE_COLLECTION_NAME,
            filter=f"knowledge_id = {knowledge_id}"
        )
        # 删除ES中的文档
        elasticsearch_object.delete_document(
            knowledge_id=knowledge_id
        )
        return {"status": "success", "message": "Knowledge deleted successfully"}
    except Exception as e:
        logger.error(f"Error deleting knowledge: {e}")
        return {"status": "error", "message": str(e)}


@app.get('/file/summary/create/page', summary="创建每一页的文章摘要")
async def file_summary(
        knowledge_id: int,
        file_hase: str
):
    """
    获取文件摘要
    """
    sql = "SELECT file_path FROM knowledge_file WHERE knowledge_id = ? AND file_key = ?"
    try:
        with sqlite_db as db_client:
            result = db_client.execute(sql, (knowledge_id, file_hase))
            if not result:
                return {"status": "error", "message": "File not found"}
            file_path = result[0][0]
        # 使用ZhipuFileUtil获取文章摘要
        vllm_file_util = VLLMFileUtil()
        file_content_list = await vllm_file_util.pdf_to_text(file_path=file_path)
        if not file_content_list:
            return {"status": "error", "message": "No content extracted from the file"}
    except Exception as e:
        logger.error(f"Error fetching file path: {e}")
        return {"status": "error", "message": str(e)}


if __name__ == '__main__':
    import uvicorn

    # 启动FastAPI应用
    uvicorn.run(app, host="0.0.0.0", port=8000)
